from aiojobs import Scheduler, Job
from typing import Dict, Callable, Any, AsyncGenerator
import asyncio

from common.log import log
from common.exception.errors import (
    ServerConflict, ServerNotFound,
    ServerError
)

from config.base import get_settings

settings = get_settings()

class AsyncTaskScheduler:
    def __init__(
        self, 
        concurrent_tasks: int = None,
        task_pending: int = None
    ):
        self._scheduler: Scheduler = None
        self._concurrent_tasks: int = concurrent_tasks or settings.task.CONCURRENCY
        self._task_pending: int = task_pending or settings.task.PENDING
        self._task_map: Dict[str, Job] = {}  # 存储任务状态
        self._lock = asyncio.Lock()
        self._shutdown_event = asyncio.Event()

    async def start(self):
        """启动调度器（在 FastAPI 生命周期中初始化）"""
        if not self._scheduler or self._scheduler.closed:
            # 设置最大并发任务数，并关闭 pending 任务自动清理
            self._scheduler = Scheduler(
                limit=self._concurrent_tasks, 
                pending_limit=self._task_pending
            )
            await self._scheduler.__aenter__()
            log.info("任务调度器已启动")

    async def stop(self):
        """停止调度器（在 FastAPI 关闭时调用）"""
        if self._scheduler and not self._scheduler.closed:
            self._shutdown_event.set()
            await self._scheduler.close()
            self._scheduler = None
            log.info("任务调度器已停止")

    async def submit_task(self, task_id: str, coro: AsyncGenerator):
        """提交任务并关联ID"""
        async with self._lock:
            await self._check_preconditions(task_id)
            
            job = await self._scheduler.spawn(self._safe_wrap_task(coro, task_id))
            self._task_map[task_id] = job
            log.info(f"任务 {task_id} 已提交，Job ID: {id(job)}")


    async def cancel_task(self, task_id: str):
        """通过任务ID取消任务"""
        async with self._lock:
            job = self._task_map.get(task_id)
            if not job:
                raise ServerNotFound("任务不存在")
            
            await job.close()  # 取消任务
            self._task_map.pop(task_id)
            log.info(f"任务 {task_id} 已取消")
            
    async def _check_preconditions(self, task_id: str) -> None:
        """提交前的条件检查"""
        if self._shutdown_event.is_set():
            raise ServerError("调度器正在关闭，拒绝新任务")
        if task_id in self._task_map:
            raise ServerConflict(f"任务ID {task_id} 已存在")
        if len(self._task_map) >= self._task_pending:
            raise ServerError("任务数达到上限")
            
    async def _safe_wrap_task(self, coro, task_id):
        """包装任务，确保状态清理"""
        try:
            await coro
            log.info(f"任务 {task_id} 执行成功")
        except asyncio.CancelledError:
            log.warning(f"任务 {task_id} 被取消")
        except Exception as e:
            log.error(
                f"任务 {task_id} 执行失败: {str(e)}",
                exc_info=True  # 记录完整堆栈
            )
        finally:
            async with self._lock:
                self._task_map.pop(task_id, None)